import csv
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.stats import norm

class SimulationConfig:
    """模拟配置类，管理所有物理参数""" 
    def __init__(self, config_dict=None):
        if config_dict is None:
            # 默认配置
            self.W_mix = 20 #0.4 * 28 + 0.6 * 41  # 混合气体平均电离能 (eV/电荷)
            self.Fano_factor = 0.3            # Fano因子
            self.GMCP_gain = 10**3            # 电子倍增因子
            self.hole_efficiency = 0.96       # 孔道入口效率
            self.collection_efficiency = 0.11 # TopMetal-II的收集率
            self.v_drift = 0.34               # 纵向漂移速度 (cm/μs) @ 1kV/cm电场强度
            self.bottom_adsorption = 0.26     # 底部吸附率
            self.MCP_Z_POSITION = -6.9        # MCP位于Z=-6.9mm平面
            self.TOPMETAL_Z_POSITION = self.MCP_Z_POSITION - 3.4  # TopMetal位于MCP下方3.4mm处
            self.DRIFT_REGION_DIFFUSION = 0.02066865  # 漂移区扩散系数 (cm/√cm)
            self.INDUCTION_REGION_E_FIELD_RATIO = 2.0  # 感应区电场强度是漂移区的2倍
            self.LONGITUDINAL_RATIO = 0.5    # 纵向扩散与横向扩散的比例
            self.SAMPLING_INTERVAL = 5       # 采样间隔 (ns)
            self.VERTICAL_DIFFUSION = 0.1    # 纵向均匀扩散系数 (mm/√ns)
        else:
            # 从字典加载配置
            self.__dict__.update(config_dict)
    
    def get_diffusion_model(self):
        """创建并返回扩散模型实例"""
        return DiffusionModel(self)

class DiffusionModel:
    """抽象扩散模型基类"""
    def __init__(self, config):
        self.config = config
    
    def calculate_drift_distance(self, position_mm, target_z_mm):
        """计算电子到目标平面的漂移距离（单位：cm）"""
        z_position_mm = position_mm[2]
        drift_distance_mm = abs(z_position_mm - target_z_mm)
        return drift_distance_mm / 10  # 转换为cm
    
    def calculate_drift_sigma(self, position):
        """计算漂移过程的扩散标准差"""
        drift_distance_cm = self.calculate_drift_distance(position, self.config.MCP_Z_POSITION)
        transverse_sigma_cm = self.config.DRIFT_REGION_DIFFUSION * np.sqrt(drift_distance_cm)
        return transverse_sigma_cm * 10  # 转换为mm
    
    def calculate_induction_sigma(self):
        """计算感应过程的扩散标准差"""
        drift_distance_cm = self.calculate_drift_distance(
            (0, 0, self.config.MCP_Z_POSITION), 
            self.config.TOPMETAL_Z_POSITION
        )
        # 扩散系数与电场强度的平方根成反比
        adjusted_diffusion = self.config.DRIFT_REGION_DIFFUSION / np.sqrt(self.config.INDUCTION_REGION_E_FIELD_RATIO)
        transverse_sigma_cm = adjusted_diffusion * np.sqrt(drift_distance_cm)
        return transverse_sigma_cm * 10  # 转换为mm
    
    def get_induction_drift_velocity(self):
        """获取感应区的漂移速度（电场强度加倍，速度加倍）"""
        return self.config.v_drift * self.config.INDUCTION_REGION_E_FIELD_RATIO
    
    def apply_diffusion(self, edep_position):
        """应用扩散模型，返回最终位置（考虑纵向扩散）"""
        drift_sigma = self.calculate_drift_sigma(edep_position)
        induction_sigma = self.calculate_induction_sigma()
        
        # 第一阶段扩散（沉积→MCP）
        mcp_x = edep_position[0] + np.random.normal(0, drift_sigma)
        mcp_y = edep_position[1] + np.random.normal(0, drift_sigma)
        mcp_z = self.config.MCP_Z_POSITION + np.random.normal(0, drift_sigma * self.config.LONGITUDINAL_RATIO)
        
        # 第二阶段扩散（MCP→Topmetal）
        topmetal_x = mcp_x + np.random.normal(0, induction_sigma)
        topmetal_y = mcp_y + np.random.normal(0, induction_sigma)
        topmetal_z = self.config.TOPMETAL_Z_POSITION + np.random.normal(0, induction_sigma * self.config.LONGITUDINAL_RATIO)
        
        # 计算横向偏移量
        delta_x_mcp = mcp_x - edep_position[0]
        delta_y_mcp = mcp_y - edep_position[1]
        delta_x_topmetal = topmetal_x - mcp_x
        delta_y_topmetal = topmetal_y - mcp_y
        
        # 最终位置（Z坐标保持为原始沉积位置，但记录纵向扩散总量）
        final_x = edep_position[0] + delta_x_mcp + delta_x_topmetal
        final_y = edep_position[1] + delta_y_mcp + delta_y_topmetal
        final_z = edep_position[2]  # 保持Z位置为原始沉积位置
        
        # 纵向扩散总量（用于统计分析）
        total_longitudinal_diffusion = (mcp_z - self.config.MCP_Z_POSITION) + (topmetal_z - self.config.TOPMETAL_Z_POSITION)
        
        return final_x, final_y, final_z, total_longitudinal_diffusion, self.calculate_arrival_time(edep_position)
    
    def calculate_arrival_time(self, edep_position):
        """计算电荷到达时间"""
        t_drift_to_mcp = abs(edep_position[2] - self.config.MCP_Z_POSITION) / (self.config.v_drift * 10)
        t_drift_to_topmetal = abs(self.config.MCP_Z_POSITION - self.config.TOPMETAL_Z_POSITION) / (self.get_induction_drift_velocity() * 10)
        return t_drift_to_mcp + t_drift_to_topmetal

def read_deposits_info(csv_path):
    """读取CSV文件中的能量沉积事件"""
    deposits = []
    with open(csv_path, 'r') as f:
        reader = csv.reader(f)
        in_deposits = False
        header = None
        for row in reader:
            row = [cell.strip() for cell in row]
            if not row:
                continue
            if "Deposits Info" in row:
                in_deposits = True
                try:
                    header = next(reader)
                    print(f"找到Deposits Info,表头字段：{header}")
                except StopIteration:
                    print("错误：存款信息部分缺少表头")
                    return []
                continue
            if in_deposits and header:
                try:
                    pre_x = float(row[header.index('pre_pos_x')])
                    pre_y = float(row[header.index('pre_pos_y')])
                    pre_z = float(row[header.index('pre_pos_z')])
                    post_x = float(row[header.index('post_pos_x')])
                    post_y = float(row[header.index('post_pos_y')])
                    post_z = float(row[header.index('post_pos_z')])
                    edep_MeV = float(row[header.index('edep')])
                    
                    deposits.append({
                        'pre_pos_x': pre_x,
                        'pre_pos_y': pre_y,
                        'pre_pos_z': pre_z,
                        'post_pos_x': post_x,
                        'post_pos_y': post_y,
                        'post_pos_z': post_z,
                        'edep_MeV': edep_MeV
                    })
                except ValueError as e:
                    print(f"数据解析错误：{e}，跳过该行数据：{row}")
    return deposits

def process_deposits(deposits, config=None):
    """处理能量沉积事件"""
    if config is None:
        config = SimulationConfig()
    
    diffusion_model = config.get_diffusion_model()
    charge_positions = []
    
    for dep in deposits:
        if dep['edep_MeV'] <= 0:
            continue
            
        mid_x = (dep['pre_pos_x'] + dep['post_pos_x']) / 2
        mid_y = (dep['pre_pos_y'] + dep['post_pos_y']) / 2
        mid_z = (dep['pre_pos_z'] + dep['post_pos_z']) / 2
        edep_position = (mid_x, mid_y, mid_z)
        
        # 计算原始电离电荷数
        edep_eV = dep['edep_MeV'] * 1e6
        mean_ion = edep_eV / config.W_mix
        sigma_ion = np.sqrt(config.Fano_factor * mean_ion)
        n_ion = max(1, int(np.round(norm.rvs(loc=mean_ion, scale=sigma_ion))))
        
        # 应用GMCP放大效应
        effective_ion = n_ion * config.hole_efficiency * (1 - config.bottom_adsorption)
        amplified_charge = int(effective_ion * config.GMCP_gain * config.collection_efficiency)
        
        # 为每个电荷应用扩散（考虑纵向扩散）
        for _ in range(amplified_charge):
            final_x, final_y, final_z, long_diff, t_total = diffusion_model.apply_diffusion(edep_position)
            charge_positions.append( (final_x, final_y, final_z, long_diff, t_total, edep_position[2]) )
    
    return np.array(charge_positions)

def sample_trajectory(charge_positions, config=None):
    """对径迹进行采样并添加纵向扩散"""
    if config is None:
        config = SimulationConfig()
    
    if len(charge_positions) == 0:
        print("错误:无电荷数据用于采样")
        return []
    
    # 提取数据
    x, y, z, long_diff, arrival_time, original_z = charge_positions.T
    
    # 找到径迹的最高点和最低点
    max_z = np.max(z)
    min_z = np.min(z)
    
    # 计算对应的时间范围（假设Z方向速度恒定）
    max_time = arrival_time[np.argmax(z)]
    min_time = arrival_time[np.argmin(z)]
    
    # 生成采样时间点（从径迹最上面开始，每隔5ns采样一次）
    sample_times = np.arange(min_time, max_time + config.SAMPLING_INTERVAL, config.SAMPLING_INTERVAL)
    
    # 为每个采样时间点创建一个采样集
    sampled_positions = []
    
    for t in sample_times:
        # 选择在该时间点附近的电荷
        mask = np.abs(arrival_time - t) < config.SAMPLING_INTERVAL / 2
        if not np.any(mask):
            continue
            
        # 提取这些电荷的位置
        sample_x = x[mask]
        sample_y = y[mask]
        sample_original_z = original_z[mask]  # 使用原始沉积位置的Z坐标
        
        # 为每个电荷添加纵向均匀扩散
        diffusion_amount = config.VERTICAL_DIFFUSION * np.sqrt(t)
        diffused_z = sample_original_z + np.random.normal(0, diffusion_amount, size=len(sample_x))
        
        # 添加到采样后的径迹中
        for i in range(len(sample_x)):
            sampled_positions.append( (sample_x[i], sample_y[i], diffused_z[i], t) )
    
    return np.array(sampled_positions)

def plot_comparison(original_positions, sampled_positions, config=None):
    """绘制采样前后的径迹对比图"""
    if config is None:
        config = SimulationConfig()
    
    if len(original_positions) == 0 or len(sampled_positions) == 0:
        print("错误:无有效电荷数据用于绘图")
        return
    
    # 提取原始径迹数据
    orig_x, orig_y, orig_z, orig_long_diff, orig_time, _ = original_positions.T
    
    # 提取采样后径迹数据
    samp_x, samp_y, samp_z, samp_time = sampled_positions.T
    
    # 创建两个图形
    fig = plt.figure(figsize=(20, 10))
    
    # 第一个子图：采样前的径迹（蓝色表示）
    ax1 = fig.add_subplot(121, projection='3d')
    sc1 = ax1.scatter(
        orig_x, orig_y, orig_z, 
        s=2,
        alpha=0.4,
        color='blue',  # 使用蓝色表示采样前的电子
        edgecolor='none',
        label='原始径迹'
    )
    
    # 绘制MCP和Topmetal平面
    x_lim = ax1.get_xlim()
    y_lim = ax1.get_ylim()
    xx, yy = np.meshgrid(x_lim, y_lim)
    ax1.plot_surface(xx, yy, np.full_like(xx, config.MCP_Z_POSITION), alpha=0.2, color='r', label='MCP')
    ax1.plot_surface(xx, yy, np.full_like(xx, config.TOPMETAL_Z_POSITION), alpha=0.2, color='b', label='Topmetal')
    
    ax1.set_xlabel('X (mm)')
    ax1.set_ylabel('Y (mm)')
    ax1.set_zlabel('Z (mm)')
    ax1.set_title(f'采样前的三维径迹\n'
                 f'[颜色: 蓝色, 纵向/横向扩散比: {config.LONGITUDINAL_RATIO}]')
    ax1.view_init(elev=15, azim=45)
    
    # 第二个子图：采样后的径迹（紫色表示）
    ax2 = fig.add_subplot(122, projection='3d')
    sc2 = ax2.scatter(
        samp_x, samp_y, samp_z, 
        s=2,
        alpha=0.4,
        color='purple',  # 使用紫色表示采样后的电子
        edgecolor='none',
        label='采样后径迹'
    )
    
    # 绘制MCP和Topmetal平面
    ax2.plot_surface(xx, yy, np.full_like(xx, config.MCP_Z_POSITION), alpha=0.2, color='r', label='MCP')
    ax2.plot_surface(xx, yy, np.full_like(xx, config.TOPMETAL_Z_POSITION), alpha=0.2, color='b', label='Topmetal')
    
    ax2.set_xlabel('X (mm)')
    ax2.set_ylabel('Y (mm)')
    ax2.set_zlabel('Z (mm)')
    ax2.set_title(f'采样后的三维径迹\n'
                 f'[颜色: 紫色, 纵向扩散系数: {config.VERTICAL_DIFFUSION}mm/√ns]')
    ax2.view_init(elev=15, azim=45)
    
    plt.tight_layout()
    plt.savefig('gmcp_trajectory_comparison.png', dpi=300, bbox_inches='tight')
    plt.show()

def validate_diffusion_results(positions, deposits, config=None):
    """验证扩散结果的统计特性（考虑纵向扩散）"""
    if config is None:
        config = SimulationConfig()
    
    if len(positions) == 0:
        print("无有效电荷数据用于验证")
        return
    
    # 计算所有沉积点的平均位置
    edep_centers = np.array([
        [(d['pre_pos_x'] + d['post_pos_x'])/2, 
         (d['pre_pos_y'] + d['post_pos_y'])/2, 
         (d['pre_pos_z'] + d['post_pos_z'])/2] 
        for d in deposits
    ])
    
    # 计算扩散后的电荷分布统计
    x, y, z, long_diff, arrival_time, _ = positions.T
    charge_mean = np.mean([x, y, z], axis=1)
    charge_std = np.std([x, y, z], axis=1)
    
    # 计算预期的横向扩散标准差
    avg_drift_distance = np.mean([
        calculate_drift_distance(center, config.MCP_Z_POSITION)
        for center in edep_centers
    ])
    
    induction_distance = calculate_drift_distance(
        (0, 0, config.MCP_Z_POSITION), 
        config.TOPMETAL_Z_POSITION
    )
    
    # 考虑电场差异的扩散系数
    drift_diffusion = config.DRIFT_REGION_DIFFUSION
    induction_diffusion = drift_diffusion / np.sqrt(config.INDUCTION_REGION_E_FIELD_RATIO)
    
    expected_transverse_sigma = np.sqrt(
        drift_diffusion**2 * avg_drift_distance +
        induction_diffusion**2 * induction_distance
    ) * 10  # 转换为mm
    
    # 计算预期的纵向扩散标准差
    expected_longitudinal_sigma = expected_transverse_sigma * config.LONGITUDINAL_RATIO
    
    # 计算实际纵向扩散统计
    long_diff_mean = np.mean(long_diff)
    long_diff_std = np.std(long_diff)
    
    print(f"电荷分布统计:")
    print(f"  平均位置: ({charge_mean[0]:.3f}, {charge_mean[1]:.3f}, {charge_mean[2]:.3f}) mm")
    print(f"  横向标准差: ({charge_std[0]:.3f}, {charge_std[1]:.3f}) mm")
    print(f"  预期横向标准差: {expected_transverse_sigma:.3f} mm")
    print(f"  纵向扩散统计:")
    print(f"    平均纵向扩散: {long_diff_mean:.3f} mm")
    print(f"    纵向扩散标准差: {long_diff_std:.3f} mm")
    print(f"    预期纵向标准差: {expected_longitudinal_sigma:.3f} mm")
    
    # 打印电场和速度信息
    induction_velocity = config.v_drift * config.INDUCTION_REGION_E_FIELD_RATIO
    print(f"\n电场和速度信息:")
    print(f"  漂移区电场强度: 1.0 (相对单位)")
    print(f"  感应区电场强度: {config.INDUCTION_REGION_E_FIELD_RATIO} (相对单位)")
    print(f"  漂移区漂移速度: {config.v_drift:.3f} cm/μs")
    print(f"  感应区漂移速度: {induction_velocity:.3f} cm/μs")

def calculate_drift_distance(position, target_z):
    """辅助函数：计算漂移距离（单位：cm）"""
    z_position = position[2]
    drift_distance_mm = abs(z_position - target_z)
    return drift_distance_mm / 10  # 转换为cm

if __name__ == "__main__":
    csv_path = "/opt/star-xp-master/build/source/0-Simulation/e-+,0.01,10/event_1.csv"
    print(f"正在读取数据文件：{csv_path}")
    
    # 读取能量沉积事件
    deposits = read_deposits_info(csv_path)
    if not deposits:
        print("错误：未读取到任何能量沉积信息")
        exit(1)
    
    # 配置模拟参数
    config = SimulationConfig()
    
    # 处理数据并应用扩散
    print("\n处理能量沉积事件...")
    charge_positions = process_deposits(deposits, config)
    print(f"生成电荷总数: {len(charge_positions)}")
    
    # 验证扩散结果（考虑纵向扩散）
    validate_diffusion_results(charge_positions, deposits, config)
    
    # 对径迹进行采样并添加纵向扩散
    print("\n对径迹进行采样并添加纵向扩散...")
    sampled_positions = sample_trajectory(charge_positions, config)
    print(f"采样后电荷总数: {len(sampled_positions)}")
    
    # 绘制对比图
    print("\n绘制采样前后的径迹对比图...")
    plot_comparison(charge_positions, sampled_positions, config)
    